import itertools
import numpy as np

__all__ = ['sort_points', 'sort_points_v2', 'sort_any_curve_points']

def sort_points(shapes, kp_num):
    point_dicts = {}
    for s in shapes:
        sk = int(s['label'])
        if sk not in point_dicts:
            point_dicts[sk] = [np.array(s['points'][0])]
        else:
            point_dicts[sk].append(np.array(s['points'][0]))

    assert len(point_dicts) == kp_num // 4

    sort_point = []
    centers = []
    for i in range(kp_num // 4):
        sps = np.array(point_dicts[i + 1])
        sps = sps[sps[:, 0].argsort()]  # x方向排序
        left_p = sps[0: 2, :]
        right_p = sps[2: 4, :]
        left_p = left_p[left_p[:, 1].argsort()]  # y方向排序
        right_p = right_p[right_p[:, 1].argsort()]  # y方向排序
        sort_point.append(left_p[0])
        sort_point.append(right_p[0])
        sort_point.append(left_p[1])
        sort_point.append(right_p[1])
        centers.append(np.mean(sps, axis=0))

    centers = np.asarray(centers, np.float32)
    sort_point = np.asarray(sort_point, np.float32)
    sort_cnt_idx = np.argsort(centers[:, 1])
    centers = centers[sort_cnt_idx]
    new_bboxes = []
    for sort_i in sort_cnt_idx:
        new_bboxes.append(sort_point[4 * sort_i, :])
        new_bboxes.append(sort_point[4 * sort_i + 1, :])
        new_bboxes.append(sort_point[4 * sort_i + 2, :])
        new_bboxes.append(sort_point[4 * sort_i + 3, :])
    new_bboxes = np.asarray(new_bboxes, np.float32)
    return new_bboxes, centers


def __is_quadrilateral(pt1:np.ndarray, pt2:np.ndarray):
    """
    判断输入的两条线段是否组成四边形
    :param pt1: [[s_x, s_y],[e_x, e_y]]
    :param pt2: [[s_x, s_y],[e_x, e_y]]
    :return: True or Flase
    """
    if __has_cross_point(pt1, pt2):
        cross_pt = __get_cross_point(pt1, pt2)
        if cross_pt[0] >= min(pt1[0][0], pt1[1][0]) and  cross_pt[0] <= max(pt1[0][0], pt1[1][0]) and \
                cross_pt[1] >= min(pt1[0][1], pt1[1][1]) and cross_pt[1] <= max(pt1[0][1], pt1[1][1]):
            return False
    return True


def __has_cross_point(pt1:np.ndarray, pt2:np.ndarray):
    """
    计算两条直线是否有相交点
    :param pt1: [[s_x, s_y],[e_x, e_y]]
    :param pt2: [[s_x, s_y],[e_x, e_y]]
    :return: True or Flase
    """
    pt1 = pt1.astype(np.float)
    pt2 = pt2.astype(np.float)
    k1 = (pt1[0][1] - pt1[1][1]) / ((pt1[0][0] - pt1[1][0]) + 0.00000001)
    k2 = (pt2[0][1] - pt2[1][1]) / ((pt2[0][0] - pt2[1][0]) + + 0.00000001)
    return k1 != k2

def __get_cross_point(pt1:np.ndarray, pt2:np.ndarray):
    """
    计算两条直线的相交点, 确保输入的两条线段是从左到右
    :param pt1: [[s_x, s_y],[e_x, e_y]]
    :param pt2: [[s_x, s_y],[e_x, e_y]]
    :return: True or Flase
    """
    pt1 = pt1.astype(np.float)
    pt2 = pt2.astype(np.float)
    k1 = (pt1[0][1] - pt1[1][1]) / ((pt1[0][0] - pt1[1][0]) + 0.00000001)
    b1 = pt1[0][1] - (k1 * pt1[0][0])

    k2 = (pt2[0][1] - pt2[1][1]) / ((pt2[0][0] - pt2[1][0]) + 0.00000001)
    b2 = pt2[0][1] - (k2 * pt2[0][0])

    x = (b1 - b2) / ((k2 - k1) + 0.00000001)
    y = k1 * x + b1

    return np.array([x, y])


def calculate_line_angle(pt1:np.ndarray, pt2:np.ndarray):
    k = (pt2[1] - pt1[1]) / ( (pt2[0] - pt1[0]) + 0.00000001)
    angle = np.arctan(k) / np.pi * 180
    return round(angle, 2)


def calculate_distance(pt1:np.ndarray, pt2:np.ndarray):
    mid_pt1 = pt1.mean(axis=0)
    mid_pt2 = pt2.mean(axis=0)
    return np.sqrt((mid_pt1[0] - mid_pt2[0])**2 + (mid_pt1[1] - mid_pt2[1])**2)


# def sort_points_v2(shapes, kp_num=68):
#     """
#     解决前后位,bending位任意弯曲情况下,对椎体的上下终板进行排序
#     从L5从底往上进行排序
#     :param shapes:
#     :param kp_num:
#     :return:
#     """
#     def sort_random_forth_points(points:np.ndarray):
#         sort_point = []
#         points = points[points[:, 0].argsort()]  # x方向排序
#         left_p = points[0: 2, :]
#         right_p = points[2: 4, :]
#         left_p = left_p[left_p[:, 1].argsort()]  # y方向排序
#         right_p = right_p[right_p[:, 1].argsort()]  # y方向排序
#         sort_point.append(left_p[0])
#         sort_point.append(right_p[0])
#         sort_point.append(left_p[1])
#         sort_point.append(right_p[1])
#         return sort_point
#
#     point_dicts = {}
#     for s in shapes:
#         sk = int(s['label'])
#         if sk not in point_dicts:
#             point_dicts[sk] = [np.array(s['points'][0])]
#         else:
#             point_dicts[sk].append(np.array(s['points'][0]))
#
#     assert len(point_dicts) == kp_num // 4
#
#     sort_points = []
#     t5_up_line = None
#
#
#     def find_quad(points:np.ndarray):
#         # 1）满足组成四边形
#         # 2）判断上边到t5_up_line的与水平线的夹角
#         # 2.1)若< 70度，判断线段s_x < e_x
#         # 2.1)若> 70度，判断线段s_y < e_y
#         # 3）上边到t5_up_line的距离 大于 下边到t5_up_line的距离
#         indexs = list(range(points.shape[0]))
#         for temp_indexs in list(itertools.permutations(indexs, 4)):
#             #print(points[temp_indexs, :])
#             s_line = points[temp_indexs[0:2], :]
#             e_line = points[temp_indexs[2:], :]
#             if __is_quadrilateral(s_line, e_line):
#                 t5_angle = calculate_line_angle(t5_up_line[0], t5_up_line[1])
#                 s_angle = calculate_line_angle(s_line[0], s_line[1])
#                 e_angle = calculate_line_angle(e_line[0], e_line[1])
#
#                 t5_s = t5_angle - s_angle
#                 t5_e = t5_angle - e_angle
#
#                 abs_angle = abs((t5_s) - (t5_e))
#                 if abs_angle > 25 or abs(t5_s) > 50 or abs(t5_e) > 50:
#                     continue
#                 else:
#                     # 判断
#                     if calculate_distance(t5_up_line, s_line) < calculate_distance(t5_up_line, e_line):
#                         continue
#                     else:
#                         if abs(s_angle) <= 80:
#                             if s_line[0][0] > s_line[1][0] or e_line[0][0] > e_line[1][0]:
#                                 continue
#                         else:
#                             if s_line[0][1] > s_line[1][1] or  e_line[0][1] > e_line[1][1]:
#                                 continue
#
#                         return np.vstack((s_line, e_line))
#         return points
#
#     for i in range(kp_num // 4, 0, -1):
#         points = np.array(point_dicts[i])
#         if i == 17:
#             sort_line = sort_random_forth_points(points)
#             sort_points.append(np.array(sort_line))
#             t5_up_line = np.array(sort_line[0:2])
#         else:
#             #寻找两条线段
#             sort_line = find_quad(points)
#             sort_points.append(sort_line)
#             t5_up_line = sort_line[0:2]
#     sort_points = np.vstack(sort_points[::-1])
#
#     return sort_points

def sort_points_v2(shapes, kp_num=68):
    """
    解决前后位,bending位任意弯曲情况下,对椎体的上下终板进行排序
    从L5从底往上进行排序
    :param shapes:
    :param kp_num:
    :return:
    """
    point_dicts = {}
    for s in shapes:
        sk = int(s['label'])
        if sk not in point_dicts:
            point_dicts[sk] = [np.array(s['points'][0])]
        else:
            point_dicts[sk].append(np.array(s['points'][0]))

    assert len(point_dicts) == kp_num // 4
    return sort_any_curve_points(point_dicts, kp_num=kp_num)


def sort_any_curve_points(point_dicts, kp_num=68):
    def sort_random_forth_points(points:np.ndarray):
        sort_point = []
        points = points[points[:, 0].argsort()]  # x方向排序
        left_p = points[0: 2, :]
        right_p = points[2: 4, :]
        left_p = left_p[left_p[:, 1].argsort()]  # y方向排序
        right_p = right_p[right_p[:, 1].argsort()]  # y方向排序
        sort_point.append(left_p[0])
        sort_point.append(right_p[0])
        sort_point.append(left_p[1])
        sort_point.append(right_p[1])
        return sort_point

    sort_points = []
    t5_up_line = None

    def find_quad(points:np.ndarray):
        # 1）满足组成四边形
        # 2）判断上边到t5_up_line的与水平线的夹角
        # 2.1)若< 70度，判断线段s_x < e_x
        # 2.1)若> 70度，判断线段s_y < e_y
        # 3）上边到t5_up_line的距离 大于 下边到t5_up_line的距离
        indexs = list(range(points.shape[0]))
        for temp_indexs in list(itertools.permutations(indexs, 4)):
            #print(points[temp_indexs, :])
            s_line = points[temp_indexs[0:2], :]
            e_line = points[temp_indexs[2:], :]
            if __is_quadrilateral(s_line, e_line):
                t5_angle = calculate_line_angle(t5_up_line[0], t5_up_line[1])
                s_angle = calculate_line_angle(s_line[0], s_line[1])
                e_angle = calculate_line_angle(e_line[0], e_line[1])

                t5_s = t5_angle - s_angle
                t5_e = t5_angle - e_angle

                abs_angle = abs((t5_s) - (t5_e))
                if abs_angle > 25 or abs(t5_s) > 50 or abs(t5_e) > 50:
                    continue
                else:
                    # 判断
                    if calculate_distance(t5_up_line, s_line) < calculate_distance(t5_up_line, e_line):
                        continue
                    else:
                        if abs(s_angle) <= 80:
                            if s_line[0][0] > s_line[1][0] or e_line[0][0] > e_line[1][0]:
                                continue
                        else:
                            if s_line[0][1] > s_line[1][1] or  e_line[0][1] > e_line[1][1]:
                                continue

                        return np.vstack((s_line, e_line))
        return points

    for i in range(kp_num // 4, 0, -1):
        points = np.array(point_dicts[i])
        if i == 17:
            sort_line = sort_random_forth_points(points)
            sort_points.append(np.array(sort_line))
            t5_up_line = np.array(sort_line[0:2])
        else:
            #寻找两条线段
            sort_line = find_quad(points)
            sort_points.append(sort_line)
            t5_up_line = sort_line[0:2]
    sort_points = np.vstack(sort_points[::-1])

    return sort_points




if __name__ == '__main__':
    import json
    from pathlib import Path
    import cv2
    from ais.image import cv_show, keypoint_to_mask
    root = Path("/home/blake/data/medical/datasets/vertebral/AIS/bending_all")
    root = Path("/home/blake/data/medical/datasets/vertebral/AIS/bending/data/all")

    for image_file in sorted(root.glob('*.jpg')):
        #image_file = Path("/home/blake/data/medical/datasets/vertebral/AIS/bending_all/yang yan1_4.jpg")
        img = cv2.imread(str(image_file))
        print(image_file)
        with open(str(image_file.with_suffix('.json')), 'r') as f:
            data_dict: dict = json.load(f)
            landmark = data_dict['shapes']
            landmark = sort_points_v2(landmark, 68)
            mask = keypoint_to_mask(landmark, img.shape[:2])

        mask:np.ndarray = mask * 255


        cv_show('mask', mask.astype(np.uint8))
        cv2.waitKey(0)